#!/usr/bin/env python3
"""
Ghidra MCP Server
为大模型提供Ghidra逆向工程能力的MCP服务器
"""

import asyncio
import requests
from typing import Any
from mcp.server import Server
from mcp.types import Tool, TextContent
import mcp.server.stdio

GHIDRA_API_BASE = "http://localhost:2233/api"

app = Server("ghidra-mcp")


@app.list_tools()
async def list_tools() -> list[Tool]:
    return [
        Tool(
            name="get_program_info",
            description="获取当前分析程序的基本信息（名称、架构、编译器、函数数量等）",
            inputSchema={"type": "object", "properties": {}, "required": []}
        ),
        Tool(
            name="list_functions",
            description="获取程序中的函数列表，包括函数地址、名称和签名",
            inputSchema={"type": "object", "properties": {}, "required": []}
        ),
        Tool(
            name="decompile_function",
            description="反编译指定地址的函数，返回C伪代码",
            inputSchema={
                "type": "object",
                "properties": {
                    "address": {"type": "string", "description": "函数地址（如'0x140001000'）"},
                    "page": {"type": "integer", "description": "页码（从0开始，默认0）"},
                    "pageSize": {"type": "integer", "description": "每页行数（默认200）"}
                },
                "required": ["address"]
            }
        ),
        Tool(
            name="get_function_info",
            description="获取指定函数的详细信息（签名、参数、大小等）",
            inputSchema={
                "type": "object",
                "properties": {
                    "address": {"type": "string", "description": "函数地址"}
                },
                "required": ["address"]
            }
        ),
        Tool(
            name="list_strings",
            description="获取程序中的字符串列表",
            inputSchema={"type": "object", "properties": {}, "required": []}
        ),
        Tool(
            name="get_cross_references",
            description="获取指定地址的交叉引用",
            inputSchema={
                "type": "object",
                "properties": {
                    "address": {"type": "string", "description": "目标地址"}
                },
                "required": ["address"]
            }
        ),
        Tool(
            name="get_function_calls",
            description="获取指定函数调用的其他函数列表",
            inputSchema={
                "type": "object",
                "properties": {
                    "address": {"type": "string", "description": "函数地址"}
                },
                "required": ["address"]
            }
        ),
        Tool(
            name="rename_function",
            description="重命名指定地址的函数",
            inputSchema={
                "type": "object",
                "properties": {
                    "address": {"type": "string", "description": "函数地址"},
                    "name": {"type": "string", "description": "新的函数名"}
                },
                "required": ["address", "name"]
            }
        ),
        Tool(
            name="add_comment",
            description="在指定地址添加注释",
            inputSchema={
                "type": "object",
                "properties": {
                    "address": {"type": "string", "description": "地址"},
                    "comment": {"type": "string", "description": "注释内容"}
                },
                "required": ["address", "comment"]
            }
        ),
        Tool(
            name="get_function_asm",
            description="获取指定函数的汇编指令列表",
            inputSchema={
                "type": "object",
                "properties": {
                    "address": {"type": "string", "description": "函数地址（如'0x140001000'）"},
                    "start": {"type": "integer", "description": "起始指令索引（从0开始）"},
                    "count": {"type": "integer", "description": "返回的指令数量（默认200）"}
                },
                "required": ["address"]
            }
        ),
        Tool(
            name="get_memory_data",
            description="获取指定地址范围的机器码数据",
            inputSchema={
                "type": "object",
                "properties": {
                    "address": {"type": "string", "description": "起始地址（如'0x140001000'）"},
                    "length": {"type": "integer", "description": "要读取的字节数（1-1024）"}
                },
                "required": ["address", "length"]
            }
        ),
        Tool(
            name="get_memory_regions",
            description="获取程序中所有内存区域的详细信息",
            inputSchema={"type": "object", "properties": {}, "required": []}
        ),
        Tool(
            name="search_memory",
            description="在内存中搜索指定的十六进制模式",
            inputSchema={
                "type": "object",
                "properties": {
                    "pattern": {"type": "string", "description": "要搜索的十六进制模式（如'A1 B2 C3'）"},
                    "start_address": {"type": "string", "description": "搜索起始地址（可选）"},
                    "end_address": {"type": "string", "description": "搜索结束地址（可选）"},
                    "limit": {"type": "integer", "description": "结果数量限制（默认100）"}
                },
                "required": ["pattern"]
            }
        ),
        Tool(
            name="write_memory",
            description="向指定地址写入数据（需谨慎操作）",
            inputSchema={
                "type": "object",
                "properties": {
                    "address": {"type": "string", "description": "目标地址"},
                    "data": {"type": "string", "description": "要写入的数据"},
                    "data_type": {"type": "string", "description": "数据类型（hex或ascii，默认hex）"}
                },
                "required": ["address", "data"]
            }
        )
    ]


@app.call_tool()
async def call_tool(name: str, arguments: Any) -> list[TextContent]:
    try:
        if name == "get_program_info":
            response = requests.get(f"{GHIDRA_API_BASE}/info", timeout=30)
            response.raise_for_status()
            data = response.json()
            result = f"""# 程序信息
**名称**: {data['name']}
**路径**: {data['executable_path']}
**架构**: {data['language']}
**编译器**: {data['compiler']}
**镜像基址**: {data['image_base']}
**内存大小**: {data['memory_size']:,} 字节
**函数数量**: {data['function_count']}"""
            return [TextContent(type="text", text=result)]
        
        elif name == "list_functions":
            response = requests.get(f"{GHIDRA_API_BASE}/functions", timeout=30)
            response.raise_for_status()
            data = response.json()
            result = f"# 函数列表（显示{data['count']}个，共{data['total']}个）\n\n"
            for func in data['functions'][:20]:
                result += f"- **{func['name']}** @ `{func['address']}`\n  ```c\n  {func['signature']}\n  ```\n\n"
            if data['count'] > 20:
                result += f"\n...还有 {data['count'] - 20} 个函数未显示"
            return [TextContent(type="text", text=result)]
        
        elif name == "decompile_function":
            address = arguments["address"]
            page = arguments.get("page", 0)
            page_size = arguments.get("pageSize", 200)
            
            response = requests.post(f"{GHIDRA_API_BASE}/decompile", json={"address": address}, timeout=30)
            response.raise_for_status()
            data = response.json()
            
            # 处理分页逻辑
            decompiled_code = data['decompiled_code']
            code_lines = decompiled_code.split('\n')
            total_lines = len(code_lines)
            
            start_index = page * page_size
            end_index = min(start_index + page_size, total_lines)
            
            # 获取当前页的代码行
            current_page_code = '\n'.join(code_lines[start_index:end_index])
            
            # 构建结果
            result = f"""# 反编译: {data['name']}
**地址**: `{data['address']}`
**签名**: `{data['signature']}`
**总代码行数**: {total_lines}
**当前显示范围**: 第 {start_index + 1}-{end_index} 行

## 反编译代码
```c
{current_page_code}
```"""
            
            # 添加下一页提示（如果有更多内容）
            if end_index < total_lines:
                next_page = page + 1
                result += f"\n\n**提示**: 还有更多内容，使用 `page={next_page}` 和 `pageSize={page_size}` 获取下一页。"
            
            return [TextContent(type="text", text=result)]
        
        elif name == "get_function_info":
            address = arguments["address"]
            response = requests.post(f"{GHIDRA_API_BASE}/function", json={"address": address}, timeout=30)
            response.raise_for_status()
            data = response.json()
            result = f"""# 函数信息: {data['name']}
**地址**: `{data['address']}`
**签名**: `{data['signature']}`
**函数体大小**: {data['body_size']} 字节
**是否为Thunk**: {'是' if data['is_thunk'] else '否'}

## 参数\n"""
            if data['parameters']:
                for param in data['parameters']:
                    result += f"- `{param['type']} {param['name']}`\n"
            else:
                result += "无参数\n"
            return [TextContent(type="text", text=result)]
        
        elif name == "list_strings":
            response = requests.get(f"{GHIDRA_API_BASE}/strings", timeout=30)
            response.raise_for_status()
            data = response.json()
            result = f"# 字符串列表（共{data['count']}个）\n\n"
            for i, s in enumerate(data['strings'][:50], 1):
                result += f"{i}. `{s['address']}`: {s['value']}\n"
            if data['count'] > 50:
                result += f"\n...还有 {data['count'] - 50} 个字符串未显示"
            return [TextContent(type="text", text=result)]
        
        elif name == "get_cross_references":
            address = arguments["address"]
            response = requests.post(f"{GHIDRA_API_BASE}/xrefs", json={"address": address}, timeout=30)
            response.raise_for_status()
            data = response.json()
            result = f"# 交叉引用到 {address}（共{data['count']}个）\n\n"
            for xref in data['xrefs']:
                result += f"- 从 `{xref['from']}` ({xref['type']})\n"
            if data['count'] == 0:
                result += "没有找到交叉引用\n"
            return [TextContent(type="text", text=result)]
        
        elif name == "get_function_calls":
            address = arguments["address"]
            response = requests.post(f"{GHIDRA_API_BASE}/calls", json={"address": address}, timeout=30)
            response.raise_for_status()
            data = response.json()
            result = f"# 函数调用（共{data['count']}个）\n\n"
            for call in data['called_functions']:
                result += f"- **{call['name']}** @ `{call['address']}`\n"
            if data['count'] == 0:
                result += "此函数不调用其他函数\n"
            return [TextContent(type="text", text=result)]
        
        elif name == "get_function_asm":
            address = arguments["address"]
            # 获取分页参数，默认值为0（从第一条开始）和200（最多返回200条）
            start = arguments.get("start", 0)
            count = arguments.get("count", 200)
            
            response = requests.post(f"{GHIDRA_API_BASE}/asm", json={"address": address}, timeout=30)
            response.raise_for_status()
            data = response.json()
            
            total_instructions = len(data['instructions'])
            # 计算实际的结束索引，确保不超出范围
            end = min(start + count, total_instructions)
            # 获取分页后的指令列表
            paginated_instructions = data['instructions'][start:end]
            
            result = f"""# 汇编指令: {data['function']}
**地址**: `{data['address']}`
**总指令数量**: {total_instructions}
**当前显示范围**: 第 {start+1}-{end} 条指令

## 指令列表
"""
            for instr in paginated_instructions:
                # 将操作数数组转换为字符串
                operands_str = " ".join(instr['operands'])
                result += f"- `{instr['address']}`: {instr['mnemonic']} {operands_str}\n"
            
            # 如果还有更多指令，提示用户如何获取下一页
            if end < total_instructions:
                result += f"\n## 更多指令\n要查看更多指令，请使用 `start={end}` 和 `count={min(count, total_instructions - end)}` 参数继续获取。"

            return [TextContent(type="text", text=result)]

        elif name == "rename_function":
            address = arguments["address"]
            new_name = arguments["name"]
            response = requests.post(f"{GHIDRA_API_BASE}/rename", json={"address": address, "name": new_name}, timeout=30)
            response.raise_for_status()
            return [TextContent(type="text", text=f"✓ 成功将函数 `{address}` 重命名为 `{new_name}`")]
        
        elif name == "add_comment":
            address = arguments["address"]
            comment = arguments["comment"]
            response = requests.post(f"{GHIDRA_API_BASE}/comment", json={"address": address, "comment": comment}, timeout=30)
            response.raise_for_status()
            return [TextContent(type="text", text=f"✓ 成功在地址 `{address}` 添加注释")]
        
        # 内存相关工具
        elif name == "get_memory_data":
            address = arguments["address"]
            length = arguments["length"]
            
            response = requests.post(f"{GHIDRA_API_BASE}/memory", json={"address": address, "length": length}, timeout=30)
            response.raise_for_status()
            data = response.json()
            
            # 格式化十六进制输出，每16字节一行
            hex_data = data["hex"]
            formatted_hex = []
            for i in range(0, len(hex_data), 32):  # 32个字符 = 16个字节
                line_hex = hex_data[i:i+32]
                line_hex_formatted = ' '.join([line_hex[j:j+2] for j in range(0, len(line_hex), 2)])
                address_line = f"{int(address, 16) + i//2:#x}" if i == 0 else f"{int(address, 16) + i//2:#x}"
                formatted_hex.append(f"`{address_line}`: {line_hex_formatted}")
            
            result = f"""# 内存数据
**地址**: `{data['address']}`
**长度**: {data['length']} 字节

## 十六进制视图
"""
            result += "\n".join(formatted_hex)
            
            result += f"""

## ASCII视图
```
{data['ascii']}
```
"""
            return [TextContent(type="text", text=result)]
        
        elif name == "get_memory_regions":
            response = requests.get(f"{GHIDRA_API_BASE}/memory/regions", timeout=30)
            response.raise_for_status()
            data = response.json()
            
            result = f"# 内存区域列表（共{data['count']}个）\n\n"
            for region in data['regions']:
                permissions = []
                if region['read']: permissions.append('R')
                if region['write']: permissions.append('W')
                if region['execute']: permissions.append('X')
                perms_str = ''.join(permissions)
                
                result += f"## {region['name']}\n"
                result += f"- **起始地址**: `{region['start_address']}`\n"
                result += f"- **结束地址**: `{region['end_address']}`\n"
                result += f"- **大小**: {region['size']:,} 字节\n"
                result += f"- **权限**: {perms_str}\n"
                result += f"- **已初始化**: {'是' if region['initialized'] else '否'}\n\n"
            
            return [TextContent(type="text", text=result)]
        
        elif name == "search_memory":
            pattern = arguments["pattern"]
            # 去除空格，确保是纯十六进制字符串
            pattern = pattern.replace(" ", "")
            
            # 构建请求参数
            request_data = {"pattern": pattern}
            if "start_address" in arguments:
                request_data["start_address"] = arguments["start_address"]
            if "end_address" in arguments:
                request_data["end_address"] = arguments["end_address"]
            if "limit" in arguments:
                request_data["limit"] = arguments["limit"]
            
            response = requests.post(f"{GHIDRA_API_BASE}/memory/search", json=request_data, timeout=30)
            response.raise_for_status()
            data = response.json()
            
            result = f"# 内存搜索结果（共{data['count']}个匹配）\n\n"
            if data['count'] > 0:
                result += "匹配地址：\n"
                for i, match in enumerate(data['results'], 1):
                    result += f"{i}. `{match['address']}`\n"
            else:
                result += "未找到匹配的模式。\n"
            
            return [TextContent(type="text", text=result)]
        
        elif name == "write_memory":
            address = arguments["address"]
            data = arguments["data"]
            data_type = arguments.get("data_type", "hex")
            
            response = requests.post(f"{GHIDRA_API_BASE}/memory/write", 
                                    json={"address": address, "data": data, "data_type": data_type}, 
                                    timeout=30)
            response.raise_for_status()
            result_data = response.json()
            
            if result_data['success']:
                return [TextContent(type="text", text=f"✓ 成功向地址 `{address}` 写入 {result_data['written_bytes']} 字节数据")]
            else:
                return [TextContent(type="text", text=f"✗ 写入内存失败")]
        
        else:
            return [TextContent(type="text", text=f"未知工具: {name}")]
            
    except requests.exceptions.HTTPError as e:
        error_msg = f"API错误 ({e.response.status_code}): "
        try:
            error_data = e.response.json()
            error_msg += error_data.get('error', str(e))
        except:
            error_msg += str(e)
        return [TextContent(type="text", text=error_msg)]
    except Exception as e:
        return [TextContent(type="text", text=f"错误: {str(e)}\n\n请确保Ghidra服务器正在运行（端口2233）")]


async def main():
    async with mcp.server.stdio.stdio_server() as (read_stream, write_stream):
        await app.run(read_stream, write_stream, app.create_initialization_options())


if __name__ == "__main__":
    def test_api():
        try:
            headers = {"User-Agent": "Mozilla/5.0"}
            resp = requests.get(f"{GHIDRA_API_BASE}/info", headers=headers, timeout=10)
            print("Status:", resp.status_code)
            print("Response:", resp.text)
        except Exception as e:
            print("Error:", e)
    
    #test_api()
    # 运行主程序
    asyncio.run(main())